"""
建立聊天类
chatbot类
0)创建工具
1)初始大模型， _init  大模型绑定工具
2)创建聊天节点 调用大模型 返回结果
3)创建工具和节点选择路由
4)初始graph  _init  使用ToolNode(tools=self._tools))工具转节点,建立条件边路由
5）运行graph __call
"""
from langgraph.graph import StateGraph,START,END,MessagesState
from model_utils import getLLM
from langgraph.checkpoint.memory import MemorySaver
import uuid
from langchain_core.tools import tool
from typing import Annotated
from app_utils import websearch,rag
from langgraph.prebuilt import ToolNode


@tool
def web_search(query:Annotated[str,"网络查询内容"]):
    '''web网络信息查询 搜索；如果网络信息查询，使用web_search'''
    return websearch(query)
@tool
def rag_search(query:Annotated[str,"人物查询字段"]):
    """如果用户查询的信息是关于人物的信息，使用chatbot_retrieve工具"""
    return rag(query)

class Chatbot:
    def __init__(self):
        self._tools = [web_search,rag_search]
        self.mermory = MemorySaver()
        self.llm = getLLM().bind_tools(self._tools)
        self.graph = self.init_graph()

    def __call__(self):
        conf = {"thread_id":uuid.uuid4()}
        while True:
            human = input("我：")
            if human == "bye":
                print("AI:bye")
                break
            r = self.graph.invoke({"messages":[human]},config=conf)
            print("AI:",r["messages"][-1].content)

    def tools_node(self,state):
        tools_info = state["messages"][-1]
        if tools_info.tool_calls:
            return "tools_node"
        return END


    def init_graph(self):
        builder = StateGraph(MessagesState)
        builder.add_node("chat_node",self.chat_node)
        builder.add_node("tools_node",ToolNode(self._tools))
        builder.add_edge(START,"chat_node")
        builder.add_conditional_edges("chat_node",self.tools_node)
        builder.add_edge("chat_node",END)
        return builder.compile(checkpointer=self.mermory)


    def chat_node(self,state):
        r = self.llm.invoke(state["messages"])
        return {"messages":[r]}

if __name__ == '__main__':
    chat = Chatbot()
    chat()